{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "attempted relative import with no known parent package",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[56], line 5\u001b[0m\n\u001b[0;32m      2\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mnn\u001b[39;00m\n\u001b[0;32m      3\u001b[0m \u001b[38;5;28;01mimport\u001b[39;00m \u001b[38;5;21;01mtorch\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mnn\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfunctional\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m \u001b[38;5;21;01mF\u001b[39;00m\n\u001b[1;32m----> 5\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mbconv\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m Bconv\n\u001b[0;32m      6\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mcfm\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m CFM\n\u001b[0;32m      7\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mfam\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m FAM\n",
      "\u001b[1;31mImportError\u001b[0m: attempted relative import with no known parent package"
     ]
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "\n",
    "from .bconv import Bconv\n",
    "from .cfm import CFM\n",
    "from .fam import FAM\n",
    "from .pdc import PDC\n",
    "from .resnet import ResNet18,init_weight\n",
    "from .rf import RF\n",
    "from .sa import SA\n",
    "from .sca import SpatialAttention, ChannelwiseAttention\n",
    "\n",
    "class MyNet(nn.Module):\n",
    "    def __init__(self):\n",
    "        super(MyNet, self).__init__()\n",
    "        self.resnet = ResNet18()\n",
    "        # init_weight(self.resnet)\n",
    "\n",
    "        self.upsample8 = nn.Upsample(scale_factor=8, mode='bilinear', align_corners=True)\n",
    "        self.rf_low = RF(64, 64)\n",
    "        self.rf_mid = RF(128, 128)\n",
    "        self.rf_high = RF(256, 256)\n",
    "        self.rf_final = RF(512, 512)\n",
    "\n",
    "        self.cfm_low = CFM(64,64,64)\n",
    "        self.cfm_high = CFM(512,256,256)\n",
    "\n",
    "        self.fam1 = FAM(512,512)\n",
    "        self.fam2 = FAM(512,256,True)\n",
    "        self.fam3 = FAM(256,128,True)\n",
    "        self.fam4 = FAM(128,64,True)\n",
    "\n",
    "        self.out1 = nn.Sequential(nn.Conv2d(64, 1, 1),FAM(1,1))\n",
    "\n",
    "        self.searchattention = SA(31,4)\n",
    "        self.cwa = ChannelwiseAttention(128)\n",
    "        self.down = nn.MaxPool2d(2, stride=2)\n",
    "\n",
    "        self.cfm_final1 = CFM(512,256,256)\n",
    "        self.cfm_final2 = CFM(256,128,128)\n",
    "\n",
    "        self.spa = SpatialAttention(128,7)\n",
    "        self.outfam = FAM(128,1)\n",
    "\n",
    "    def forward(self,x):\n",
    "        resnet_dict = self.resnet(x)\n",
    "        x0 = resnet_dict['X0']      # 64 88 88\n",
    "        x1 = resnet_dict['X1']      # 64 88 88\n",
    "        x2 = resnet_dict['X2']      # 128 44 44\n",
    "        x3 = resnet_dict['X3']      # 256 22 22\n",
    "        x4 = resnet_dict['X4']      # 512 11 11\n",
    "\n",
    "        x_low = self.cfm_low(x0, x1)\n",
    "        x_low = self.rf_low(x_low)  # 64 88 88\n",
    "\n",
    "        x_mid = self.rf_mid(x2)    # 128 44 44\n",
    "\n",
    "        x_high = self.cfm_high(x3,x4)\n",
    "        x_high = self.rf_high(x_high)   # 256 22 22\n",
    "\n",
    "        x_final = self.rf_final(x4)     # 512 11 11\n",
    "        x_b1 = self.fam1(x_final)\n",
    "        x_b2 = self.fam2(x_b1, x_high)\n",
    "        x_b3 = self.fam3(x_b2, x_mid)\n",
    "        x_b4 = self.fam4(x_b3, x_low)\n",
    "\n",
    "        out1 = self.out1(x_b4)\n",
    "        out1 = self.down(out1)\n",
    "\n",
    "        x1,_ = self.cwa(x2)\n",
    "        x5 = self.searchattention(out1.sigmoid(), x1)\n",
    "\n",
    "        x3_2 = self.resnet.layer3_im(x5)\n",
    "        x4_2 = self.resnet.layer4_im(x3_2)\n",
    "        \n",
    "        x_final1 = self.cfm_final1(x3_2, x4_2)\n",
    "        x_final2 = self.cfm_final2(x5,x_final1)\n",
    "\n",
    "        out2 = self.upsample8(x_final2)\n",
    "        out2 = self.spa(out2)\n",
    "        out2 = self.outfam(out2)\n",
    "\n",
    "        out1 = self.upsample8(out1)\n",
    "        return out1, out2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1, 1, 352, 352])\n"
     ]
    }
   ],
   "source": [
    "model = MyNet()\n",
    "input = torch.randn(1,3,352,352)\n",
    "output,_ = model(input)\n",
    "print(output.shape)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "vir2",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
